package com.yujiahao;

import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.streaming.connectors.kinesis.FlinkKinesisConsumer;
import org.apache.flink.streaming.connectors.kinesis.FlinkKinesisProducer;
import org.apache.flink.streaming.connectors.kinesis.config.AWSConfigConstants;
import org.apache.flink.streaming.connectors.kinesis.config.ConsumerConfigConstants;
import org.apache.flink.streaming.connectors.kinesis.serialization.KinesisDeserializationSchemaWrapper;
import org.apache.flink.streaming.connectors.kinesis.serialization.KinesisSerializationSchema;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Properties;

/**
 * @Description: 操作Kinesis的工具类
 * @Author: yujiahao
 * @Date: 2023/6/14 19:31
 * @Version: 1.0
 */

public class KinesisUtil {
    // AWS 访问密钥
    private static final String AWS_ACCESS_KEY_ID = "替换成自己的";
    private static final String AWS_SECRET_ACCESS_KEY = "替换成自己的";
    private static final String AWS_S3_REGION_NAME = "us-east-2";


    /**
     * TODO 1、从Kinesis数据流读取数据作为数据源
     *
     * @param streamName
     * @return
     */
    public static FlinkKinesisConsumer<String> getKinesisSource(String streamName) {

        // 配置 Amazon Kinesis 连接信息
        Properties kinesisConfig = new Properties();
        kinesisConfig.setProperty(ConsumerConfigConstants.AWS_REGION, AWS_S3_REGION_NAME);
        kinesisConfig.setProperty(ConsumerConfigConstants.STREAM_INITIAL_POSITION, "LATEST");
        kinesisConfig.setProperty(ConsumerConfigConstants.AWS_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID);
        kinesisConfig.setProperty(ConsumerConfigConstants.AWS_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY);
       //设置消费频率
        kinesisConfig.setProperty(ConsumerConfigConstants.SHARD_GETRECORDS_INTERVAL_MILLIS, "5000");
        //注册流消费者
//        kinesisConfig.put(ConsumerConfigConstants.RECORD_PUBLISHER_TYPE,
//                ConsumerConfigConstants.RecordPublisherType.EFO.name());
//        kinesisConfig.put(ConsumerConfigConstants.EFO_CONSUMER_NAME, EfoConsumerName);
//
//        kinesisConfig.put(ConsumerConfigConstants.EFO_REGISTRATION_TYPE,
//                ConsumerConfigConstants.EFORegistrationType.LAZY.name());

        FlinkKinesisConsumer<String> kinesisConsumer = new FlinkKinesisConsumer<>(streamName,
                new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
                kinesisConfig);

        return kinesisConsumer;
    }


    /**
     * TODO 2.txt、定义一个FlinkKinesisProducer，Kinesis 生产者
     * @param streamName 写入的数据流名称
     * @return
     * @param <T>
     * @throws Exception
     */
    public static <T> FlinkKinesisProducer<T> sendToKinesis(String streamName) throws Exception {
        // 配置 Kinesis 生产者属性
        Properties producerConfig = new Properties();
        producerConfig.setProperty(AWSConfigConstants.AWS_REGION, AWS_S3_REGION_NAME);
        producerConfig.setProperty(AWSConfigConstants.AWS_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID);
        producerConfig.setProperty(AWSConfigConstants.AWS_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY);
        producerConfig.put("AggregationMaxCount", "4294967295");
        producerConfig.put("RecordTtl", "30000");
        producerConfig.put("RequestTimeout", "6000");
        producerConfig.put("ThreadPoolSize", "15");

        // 创建一个带有 Kinesis 生产者的 FlinkKinesisProducer
        KinesisSerializationSchema<T> serializationSchema = new KinesisSerializationSchema<T>() {
            @Override
            public ByteBuffer serialize(T element) {
                // 序列化成字节数组
                return ByteBuffer.wrap(element.toString().getBytes(StandardCharsets.UTF_8));
            }
            @Override
            public String getTargetStream(T element) {
                return streamName;
            }
        };

        FlinkKinesisProducer<T> flinkKinesisProducer = new FlinkKinesisProducer<>(serializationSchema, producerConfig);
        flinkKinesisProducer.setDefaultStream(streamName);
        flinkKinesisProducer.setDefaultPartition("0");
        flinkKinesisProducer.setFailOnError(true);

        return flinkKinesisProducer;
    }


}
